#include <ros/ros.h>
#include <nav_msgs/Path.h>
#include <nav_msgs/OccupancyGrid.h>
#include <geometry_msgs/PoseStamped.h>
#include <geometry_msgs/PoseWithCovarianceStamped.h>
#include <tf2/utils.h>
#include <mutex>
#include "hybrid_a_star.h"

class HybridAStarROS {
    ros::NodeHandle nh_;
    ros::Subscriber map_sub_, start_sub_, goal_sub_;
    ros::Publisher path_pub_;
    HybridAStar::HybridAStar planner_;
    HybridAStar::OccupancyGrid map_;
    std::mutex mutex_;
    geometry_msgs::PoseStamped start_pose_, goal_pose_;
    bool has_start_ = false, has_goal_ = false;

    void mapCallback(const nav_msgs::OccupancyGrid::ConstPtr& msg) {
        std::lock_guard<std::mutex> lock(mutex_);
        HybridAStar::OccupancyGrid::MetaInfo meta;
        meta.width = msg->info.width;
        meta.height = msg->info.height;
        meta.resolution = msg->info.resolution;
        meta.origin_x = msg->info.origin.position.x;
        meta.origin_y = msg->info.origin.position.y;
        
        map_.load(meta, msg->data);
        planner_.setMapResolution(meta.resolution);
        ROS_INFO("Map updated | Res: %.3fm | Size: %dx%d", 
                meta.resolution, meta.width, meta.height);
    }

    void startCallback(const geometry_msgs::PoseWithCovarianceStamped::ConstPtr& msg) {
        std::lock_guard<std::mutex> lock(mutex_);
        start_pose_.header = msg->header;
        start_pose_.pose = msg->pose.pose;
        has_start_ = true;
        triggerPlanning();
    }

    void goalCallback(const geometry_msgs::PoseStamped::ConstPtr& msg) {
        std::lock_guard<std::mutex> lock(mutex_);
        goal_pose_ = *msg;
        has_goal_ = true;
        triggerPlanning();
    }

    void triggerPlanning() {
        if (!has_start_ || !has_goal_ || map_.data.empty()) return;

        HybridAStar::State start(
            start_pose_.pose.position.x,
            start_pose_.pose.position.y,
            tf2::getYaw(start_pose_.pose.orientation)
        );
        
        HybridAStar::State goal(
            goal_pose_.pose.position.x,
            goal_pose_.pose.position.y,
            tf2::getYaw(goal_pose_.pose.orientation)
        );

        ROS_INFO("Start planning...");
        auto path = planner_.plan(start, goal, map_);
        
        if (!path.empty()) {
            ROS_INFO("Found path with %zu nodes", path.size());
            publishPath(path);
        } else {
            ROS_WARN("Planning failed!");
        }
    }

    void publishPath(const std::vector<HybridAStar::State>& path) {
        nav_msgs::Path msg;
        msg.header.frame_id = "map";
        msg.header.stamp = ros::Time::now();

        for (const auto& s : path) {
            geometry_msgs::PoseStamped pose;
            pose.pose.position.x = s.x;
            pose.pose.position.y = s.y;
            tf2::Quaternion q;
            q.setRPY(0, 0, s.theta);
            tf2::convert(q, pose.pose.orientation);
            msg.poses.push_back(pose);
        }
        
        path_pub_.publish(msg);
    }

public:
    HybridAStarROS() {
        map_sub_ = nh_.subscribe("/map", 1, &HybridAStarROS::mapCallback, this);
        start_sub_ = nh_.subscribe("/initialpose", 1, &HybridAStarROS::startCallback, this);
        goal_sub_ = nh_.subscribe("/move_base_simple/goal", 1, &HybridAStarROS::goalCallback, this);
        path_pub_ = nh_.advertise<nav_msgs::Path>("/hybrid_path", 1);

        HybridAStar::VehicleParams params;
        nh_.param("vehicle_length", params.length, 4.8);
        nh_.param("vehicle_width", params.width, 2.0);
        nh_.param("wheelbase", params.wheelbase, 2.8);
        nh_.param("front_overhang", params.front_overhang, 0.9);
        nh_.param("rear_overhang", params.rear_overhang, 1.1);
        nh_.param("max_steer", params.max_steer, M_PI/4);
        
        planner_ = HybridAStar::HybridAStar(params);
        ROS_INFO("Hybrid A* initialized | Vehicle L:%.1fm W:%.1fm", 
                params.length, params.width);
    }
};

int main(int argc, char** argv) {
    ros::init(argc, argv, "hybrid_planner");
    HybridAStarROS node;
    ros::spin();
    return 0;
}